import os
import math
import logging
import pandas as pd
import psycopg2.extras
import xml.etree.ElementTree as ET
from Units import QUnits
from datetime import datetime
from pytdx.reader import BlockReader

class DrawLineUnit:
    def __init__(self, code, cycle):
        self.unit = QUnits()
        self.setting = self.unit.XML2ObjFromFile('./setting.xml').setting
        #self.filePath = self.setting.drawLine.path + '/' + self._getStockPath(code) + '/' + self._getCycleFile(cycle)
        #print(self.unit.XML2DictFromFile(self.filePath)['DrawLineData']['DrawLine'][6])
        #print(self._drawLineData()[5])
        #print(self._getDrawLineTypeIndex('name', '直线'))
        filePath = self.setting.drawLine.path + '/' + self._getStockInfo(code)[0] + '/' + self._getCycleFile(cycle)
        tree = ET.parse(filePath)
        ele = tree.findall("DrawLine[@LineType='87']")
        for r in ele:
            print(ET.tostring(r))
            r.set('IsCrossPeriod', '1')
            print(ET.tostring(r))

    def _getYData(self, price, timestamp):
        if(len(str(price).split('.')[1]) == 2):
            yData = str(price) + '000' + str(timestamp)
        elif(len(str(price).split('.')[1]) == 3):
            yData = str(price) + '00' + str(timestamp)
        else:
            yData = str(price) + '000000000'
        return yData

    def _getOpenPrice(self, cycle, market, code, start=0):
        if(cycle == '1m'):
            category = 8
        elif (cycle == '5m'):
            category = 0
        elif (cycle == '15m'):
            category = 1
        elif (cycle == '30m'):
            category = 2
        elif (cycle == '6m'):
            category = 3
        elif (cycle == '1d'):
            category = 4
        else:
            category = 5
        return self.unit.getStockSecurityBars(category, market, code, start, 1)[0]['open']

    def _getXtime(self, cycle, time=None):
        if(time == None):
            timestamp = datetime.timestamp(datetime.now())
        else:
            if (isinstance(time, datetime)):
                timestamp = datetime.timestamp(time)
            elif (isinstance(time, str)):
                dateTimeObject = datetime.strptime(time, "%Y-%m-%d %H:%M:%S")
                timestamp = datetime.timestamp(dateTimeObject)
            elif (isinstance(time, int)):
                timestamp = str(time).split('.')[0]
                if (len(timestamp) == 10):
                    timestamp = int(timestamp) / 1000
                elif (len(timestamp) == 13):
                    timestamp = int(timestamp)
                else :
                    timestamp = datetime.timestamp(datetime.now())
            else:
                timestamp = datetime.timestamp(datetime.now())
        if(cycle == 'cycle5m'):
            newTimestamp = math.ceil(int(timestamp) / (5 * 60)) * 5 * 60
            return datetime.fromtimestamp(newTimestamp).strftime("%Y%m%d%H%M") + '0000'
        elif (cycle == 'cycle15m'):
            newTimestamp = math.ceil(int(timestamp) / (15 * 60)) * 15 * 60
            return datetime.fromtimestamp(newTimestamp).strftime("%Y%m%d%H%M") + '0000'
        elif (cycle == 'cycle30m'):
            newTimestamp = math.ceil(int(timestamp) / (30 * 60)) * 30 * 60
            return datetime.fromtimestamp(newTimestamp).strftime("%Y%m%d%H%M") + '0000'
        elif (cycle == 'cycle60m'):
            newTimestamp = math.ceil(int(timestamp) / (60 * 60)) * 60 * 60
            return datetime.fromtimestamp(newTimestamp).strftime("%Y%m%d%H%M") + '0000'
        elif (cycle == 'cycle1d'):
            return datetime.fromtimestamp(timestamp).strftime("%Y%m%d") + '00000000'
        elif (cycle == 'cycle1w'):
            pass
        else:
            newTimestamp = math.ceil(int(timestamp) / 60) * 60
            return datetime.fromtimestamp(newTimestamp).strftime("%Y%m%d%H%M") + '0000'

    def _getDrawLineTypeIndex(self, kind, value):
        default = [{
            'Index': 0,
            'LineType': '87',
            'LineName': '度量目标线'
        }, {
            'Index': 1,
            'LineType': '48',
            'LineName': '射线'
        }, {
            'Index': 2,
            'LineType': '49',
            'LineName': '直线'
        }, {
            'Index': 3,
            'LineType': '16',
            'LineName': '线段'
        }, {
            'Index': 4,
            'LineType': '54',
            'LineName': '水平线'
        }, {
            'Index': 5,
            'LineType': '67',
            'LineName': '文本'
        }, {
            'Index': 6,
            'LineType': '68',
            'LineName': '注释'
        }]
        if(kind == 'name'):
            return pd.DataFrame(default).query("LineName == '%s'" % value)['Index'].values[0]
        else:
            return pd.DataFrame(default).query("LineType == '%s'" % value)['Index'].values[0]

    def _drawLineData(self):
        return [
            {
                '@LineType': '87',
                '@LineName': '度量目标线',
                '@IsCrossPeriod': '0',
                '@Width': '1',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Cyan9',
                '@LineOpacity': '0.76',
                '@DWStyle': '0',
                '@BKShow': '1',
                '@BKColor': 'Brush.Yellow6',
                '@BKOpacity': '0.15',
                '@Show': '1',
                '@MarkShow': '1',
                '@ExtLine': '4',
                'CustomConfig': [{
                    '@Index': '1',
                    '@IsChecked': '1',
                    '@Scale': '1',
                    '@BKColor': 'Brush.Green6',
                    '@BKOpacity': '0.15'
                }, {
                    '@Index': '2',
                    '@IsChecked': '1',
                    '@Scale': '2',
                    '@BKColor': 'Brush.Cyan6',
                    '@BKOpacity': '0.15'
                }],
                'BasePoint': {
                    '@XTime': '20230821134500',
                    '@YData': '-1',
                    '@OpenPrice': '15.79'
                },
                'KeyPoint': [{
                    '@XTime': '20230818135500',
                    '@Distance': '-46',
                    '@YDataType': '20480',
                    '@YData': '15.7',
                    '@OpenPrice': '15.79'
                }, {
                    '@XTime': '20230818145000',
                    '@Distance': '-35',
                    '@YDataType': '20480',
                    '@YData': '15.85',
                    '@OpenPrice': '15.83'
                }]
            }, {
                '@LineType': '48',
                '@LineName': '射线',
                '@IsCrossPeriod': '0',
                '@Width': '1',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Green6',
                '@LineOpacity': '1',
                '@DWStyle': '1',
                '@BKShow': '1',
                '@BKColor': 'Brush.Yellow6',
                '@BKOpacity': '0.15',
                '@Show': '1',
                '@ExtLine': '0',
                'BasePoint': {
                    '@XTime': '20230821134500',
                    '@YData': '-1',
                    '@OpenPrice': '15.79'
                }, 'KeyPoint': [{
                    '@XTime': '20230817103000',
                    '@Distance': '-117',
                    '@YDataType': '20480',
                    '@YData': '16.10024546107695',
                    '@OpenPrice': '16.11'
                }, {
                    '@XTime': '20230817150000',
                    '@Distance': '-81',
                    '@YDataType': '20480',
                    '@YData': '16.10024546107695',
                    '@OpenPrice': '16.08'
                }]
            }, {
                '@LineType': '49',
                '@LineName': '直线',
                '@IsCrossPeriod': '0',
                '@Width': '1',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Yellow6',
                '@LineOpacity': '1',
                '@DWStyle': '0',
                '@BKShow': '1',
                '@BKColor': 'Brush.Yellow6',
                '@BKOpacity': '0.15',
                '@Show': '1',
                '@ExtLine': '0',
                'BasePoint': {
                    '@XTime': '20230821134500',
                    '@YData': '-1',
                    '@OpenPrice': '15.79'
                },
                'KeyPoint': [{
                    '@XTime': '20230817093500',
                    '@Distance': '-128',
                    '@YDataType': '20480',
                    '@YData': '15.733045266419287',
                    '@OpenPrice': '15.74'
                }, {
                    '@XTime': '20230818142500',
                    '@Distance': '-40',
                    '@YDataType': '20480',
                    '@YData': '15.736594118016317',
                    '@OpenPrice': '15.74'
                }]
            }, {
                '@LineType': '16',
                '@LineName': '线段',
                '@IsCrossPeriod': '0',
                '@Width': '1',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Blue6',
                '@LineOpacity': '1',
                '@DWStyle': '1',
                '@BKShow': '1',
                '@BKColor': 'Brush.Yellow6',
                '@BKOpacity': '0.15',
                '@Show': '1',
                '@ExtLine': '4',
                'BasePoint': {
                    '@XTime': '20230821134500',
                    '@YData': '-1',
                    '@OpenPrice': '15.79'
                },
                'KeyPoint': [{
                    '@XTime': '20230815134000',
                    '@Distance': '-193',
                    '@YDataType': '20480',
                    '@YData': '15.61077530422818',
                    '@OpenPrice': '15.65'
                }, {
                    '@XTime': '20230821100500',
                    '@Distance': '-26',
                    '@YDataType': '20480',
                    '@YData': '15.641003223627125',
                    '@OpenPrice': '15.67'
                }]
            }, {
                '@LineType': '54',
                '@LineName': '水平线',
                '@IsCrossPeriod': '0',
                '@Width': '1',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Blue6',
                '@LineOpacity': '1',
                '@DWStyle': '1',
                '@BKShow': '0',
                '@BKColor': '',
                '@BKOpacity': '0',
                '@Show': '1',
                '@ExtLine': '0',
                'BasePoint': {
                    '@XTime': '20230821145500',
                    '@YData': '-1',
                    '@OpenPrice': '15.63'
                },
                'KeyPoint': {
                    '@XTime': '20230814111000',
                    '@Distance': '-267',
                    '@YDataType': '20480',
                    '@YData': '15.986664530679661',
                    '@OpenPrice': '15.71'
                }
            }, {
                '@LineType': '67',
                '@LineName': '文本',
                '@IsCrossPeriod': '0',
                '@Width': '2',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Yellow6',
                '@LineOpacity': '1',
                '@DWStyle': '0',
                '@BKShow': '1',
                '@BKColor': 'Brush.F1',
                '@BKOpacity': '0.15',
                '@Show': '1',
                '@ExtLine': '0',
                '@WordName': 'Microsoft Yahei, simsun, Eurostile',
                '@WordClr': 'Brush.Yellow6',
                '@TextContent': '买入，单价：16.03，数量：2000',
                '@IsBold': '0',
                '@IsUnderline': '0',
                '@WordSize': '10',
                '@BorderShow': '0',
                'BasePoint': {
                    '@XTime': '20230821145500',
                    '@YData': '-1',
                    '@OpenPrice': '15.63'
                },
                'KeyPoint': {
                    '@XTime': '20230817112000',
                    '@Distance': '-121',
                    '@YDataType': '20480',
                    '@YData': '16.101033066429448',
                    '@OpenPrice': '16.03'
                }
            }, {
                '@LineType': '68',
                '@LineName': '注释',
                '@IsCrossPeriod': '0',
                '@Width': '2',
                '@IsLocked': '0',
                '@LineColor': 'Brush.Blue6',
                '@LineOpacity': '1',
                '@DWStyle': '0',
                '@BKShow': '1',
                '@BKColor': 'Brush.Blue6',
                '@BKOpacity': '1',
                '@Show': '1',
                '@ExtLine': '0',
                '@WordName': 'Microsoft Yahei, simsun, Eurostile',
                '@WordClr': 'Brush.Gray9',
                '@TextContent': '卖出，单价：15.96，数量：2000',
                '@IsBold': '0',
                '@IsUnderline': '0',
                '@WordSize': '12',
                '@BorderShow': '0',
                'BasePoint': {
                    '@XTime': '20230821145500',
                    '@YData': '-1',
                    '@OpenPrice': '15.63'
                },
                'KeyPoint': [{
                    '@XTime': '20230816111500',
                    '@Distance': '-170',
                    '@YDataType': '20480',
                    '@YData': '15.981013920665397',
                    '@OpenPrice': '15.95'
                }, {
                    '@XTime': '20230816131000',
                    '@Distance': '-165',
                    '@YDataType': '20480',
                    '@YData': '16.072105374886316',
                    '@OpenPrice': '15.89'
                }]
            }
        ]

    def _getCycleFile(self, cycle):
        cycleDict = {
            'cycle1m': '12288.xml',
            'cycle5m': '12293.xml',
            'cycle15m': '12303.xml',
            'cycle30m': '12318.xml',
            'cycle60m': '12348.xml',
            'cycle1d': '16384.xml',
            'cycle1w': '20480.xml',
        }
        return cycleDict['cycle' + cycle]

    def _getStockInfo(self, code):
        conn = self.unit.dbConnect(self.setting.pgsql)
        try:
            typeDict = {
                'a.sh.stock': 'USHA',
                'a.sz.stock': 'USZA',
                'sme.sz.stock': 'USZA',
                'kcb.sh.stock': 'USHA',
                'gem.sz.stock': 'USZA',
                'etf.sz.fund': 'USZJ',
                'etf.sh.fund': 'USHJ'
            }
            cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
            cursor.execute("SELECT market, stock_securities FROM trading.stock_view WHERE stock_code = '%s'" % code)
            result = dict(cursor.fetchone())
            return (typeDict[result['stock_securities']] + '_' + code, result['market'])
        except Exception as e:
            logging.warning(e)
        finally:
            self.unit.dbDisconnect(conn)

    def gruckDrawLineAbstracts(self, cycle, code, operate, price, number, time=None, losses=None):
        stock = self._getStockInfo(code)
        xTime = self._getXtime('cycle' + cycle)
        if (time is None):
            timestamp = datetime.timestamp(datetime.now())
        else:
            timestamp = time
        if (losses is None):
            abstracts = operate + '，单价：' + str(price) + '，数量：' + number
            if (operate == 'sell'):
                priceTh2 = price * 1.01
            else:
                priceTh2 = price * 0.99
            openTh2 = self._getOpenPrice(cycle, stock[1], code, 5)
            yData = self._getYData(priceTh2, timestamp)
        else:
            abstracts = operate + '，数量：' + number
            openTh2 = self._getOpenPrice(cycle, stock[1], code, 10)
            yData = self._getYData(price, timestamp)

    def gruckDrawLine54(self, cycle, code, operate=None, price=None):
        filePath = self.setting.drawLine.path + '/' + self._getStockInfo(code)[0] + '/' + self._getCycleFile(cycle)
        dlpd = self.unit.Dict2XML(self._drawLineData()[4])
        print(dlpd)

if __name__ == '__main__':
    BlockReader_TYPE_FLAT = 0
    BlockReader_TYPE_GROUP = 1
    df = BlockReader().get_df("D:/通达信金融终端(开心果整合版)V2023.04/T0002/hq_cache/block_zs.dat", result_type=BlockReader_TYPE_GROUP)
    print(df)
    #print(df.loc[df["blockname"]=='沪深300', :])